#pragma once
#ifndef GROUP_GEMM_H
#define GROUP_GEMM_H

#include "aclnn_index_group_matmul.h"
#include "aclnnop/aclnn_grouped_matmul_v4.h"
#include "base/utils_print.h"
#include "base/utils_tensor.h"
namespace group_gemm {
class GroupGemmBase {
 public:
  GroupGemmBase(int32_t m,
                int32_t n,
                int32_t k,
                int32_t group_size,
                int32_t weights_per_token) {
    params.m = m;
    params.n = n;
    params.k = k;
    params.group_size = group_size;
    params.weights_per_token = weights_per_token;
    shapes.x_shape = {{params.m * 8, params.k}};
    shapes.weight_shape = {{params.group_size, params.n, params.k}};
    shapes.scale_shape = {{params.group_size, params.n}};
    shapes.per_token_scale_shape = {{params.m * 8}};
    shapes.group_list_shape = {{params.group_size}};
    shapes.y_shape = {{params.m * 8, params.n}};
  }
  void create_torch_tensors() {
    input_a = torch::randint(0, 3, {params.m * 8, params.k}, torch::kInt8);
    input_b = torch::randint(
        0, 3, {params.group_size, params.n, params.k}, torch::kInt8);
    scales = torch::rand({params.group_size, params.n}, torch::kBFloat16);
    per_token_scales = torch::rand({params.m * 8}, torch::kFloat);
    group_offset_list = GroupGemmBase::generate_groupoffset(
        params.m, params.group_size, params.weights_per_token);
    group_offset = torch::from_blob(
        group_offset_list.data(), {params.group_size}, torch::kInt64);
    output_grouped = torch::zeros({params.m * 8, params.n}, torch::kBFloat16);
    output_index = torch::zeros({params.m * 8, params.n}, torch::kBFloat16);
  }
  torch::Tensor output_grouped;
  torch::Tensor output_index;

 private:
  struct TestParams {
    int64_t m, n, k;
    int64_t group_size;
    int64_t weights_per_token;
  };

  struct TensorShapes {
    std::vector<std::vector<int64_t>> x_shape;
    std::vector<std::vector<int64_t>> weight_shape;
    std::vector<std::vector<int64_t>> scale_shape;
    std::vector<std::vector<int64_t>> per_token_scale_shape;
    std::vector<std::vector<int64_t>> group_list_shape;
    std::vector<std::vector<int64_t>> sorted_list_shape;
    std::vector<std::vector<int64_t>> y_shape;
  };

  static std::vector<int64_t> generate_groupoffset(int64_t m,
                                                   int64_t group_size,
                                                   int64_t weights_per_token) {
    std::vector<int64_t> sorted_list(m * weights_per_token);
    for (int64_t i = 0; i < m; ++i) {
      std::vector<int64_t> choices(group_size);
      std::iota(choices.begin(), choices.end(), 0);
      std::random_shuffle(choices.begin(), choices.end());
      std::copy(choices.begin(),
                choices.begin() + weights_per_token,
                sorted_list.begin() + i * weights_per_token);
    }

    std::vector<int64_t> s(group_size, 0);
    for (const auto& idx : sorted_list) {
      s[idx]++;
    }

    std::vector<int64_t> group_offset(s.size());
    std::partial_sum(s.begin(), s.end(), group_offset.begin());

    return group_offset;
  }

  int process_weight_tensor(
      const std::vector<std::vector<int64_t>>& weight_shape,
      torch::Tensor& input_b,
      aclTensorList** weight) {
    input_b = input_b.transpose(1, 2).contiguous();
    auto weight_tmp = input_b;
    weight_tmp = at_npu::native::npu_format_cast(weight_tmp.to("npu"), 29);
    std::vector<int64_t> trans_shape = {
        weight_shape[0][0], weight_shape[0][2], weight_shape[0][1]};
    int8_t* data1 = weight_tmp.data_ptr<int8_t>();
    std::vector<int64_t> strides(trans_shape.size(), 1);
    for (int64_t i = trans_shape.size() - 2; i >= 0; i--) {
      strides[i] = trans_shape[i + 1] * strides[i + 1];
    }
    std::vector<int64_t> storageShape =
        utils::get_weight_storage_shape(trans_shape);

    aclTensor* weight_nz = aclCreateTensor(trans_shape.data(),
                                           trans_shape.size(),
                                           aclDataType::ACL_INT8,
                                           strides.data(),
                                           0,
                                           aclFormat::ACL_FORMAT_FRACTAL_NZ,
                                           storageShape.data(),
                                           storageShape.size(),
                                           data1);
    std::vector<aclTensor*> tmp{weight_nz};
    *weight = aclCreateTensorList(tmp.data(), tmp.size());

    return ACL_SUCCESS;
  }

  TestParams params;
  TensorShapes shapes;
  torch::Tensor input_a;
  torch::Tensor input_b;
  torch::Tensor scales;
  torch::Tensor per_token_scales;
  torch::Tensor group_offset;
  std::vector<int64_t> group_offset_list;

  friend class GroupGemmIndx;
  friend class GroupGemmNative;
};

class GroupGemmIndx {
 public:
  GroupGemmIndx(GroupGemmBase& base) : base(base) {}
  void process(aclrtStream stream) {
    this->createTensors();
    CHECK_ACL_SUCCESS(
        this->execute_group_gemm_operator(
            x, weight, scale, perTokenScale, groupedList, y_index, stream),
        "executeMatmulOperator failed");
  }
  void destroyTensors() {
    GroupGemmIndx::clean_up(x, weight, scale, perTokenScale, y_index);
  }

 private:
  void createTensors() {
    CHECK_ACL_SUCCESS(utils::create_tensor_list_from_torch(
                          this->base.shapes.x_shape[0], this->base.input_a, &x),
                      "create x TensorList failed");
    CHECK_ACL_SUCCESS(
        utils::create_tensor_list_from_torch(
            this->base.shapes.scale_shape[0], this->base.scales, &scale),
        "create scale TensorList failed");
    CHECK_ACL_SUCCESS(utils::create_tensor_list_from_torch(
                          this->base.shapes.per_token_scale_shape[0],
                          this->base.per_token_scales,
                          &perTokenScale),
                      "create perTokenScale TensorList failed");
    CHECK_ACL_SUCCESS(
        utils::create_tensor_from_torch(this->base.shapes.group_list_shape[0],
                                        this->base.group_offset,
                                        &groupedList),
        "create groupedList Tensor failed");
    CHECK_ACL_SUCCESS(
        this->base.process_weight_tensor(
            this->base.shapes.weight_shape, this->base.input_b, &weight),
        "process weight tensor failed");
    CHECK_ACL_SUCCESS(
        utils::create_tensor_list_from_torch(
            this->base.shapes.y_shape[0], this->base.output_index, &y_index),
        "create y_index TensorList failed");
  }
  static void clean_up(aclTensorList* x,
                       aclTensorList* weight,
                       aclTensorList* scale,
                       aclTensorList* perTokenScale,
                       aclTensorList* y) {
    aclDestroyTensorList(x);
    aclDestroyTensorList(weight);
    aclDestroyTensorList(scale);
    aclDestroyTensorList(perTokenScale);
    aclDestroyTensorList(y);
  }
  int execute_group_gemm_operator(aclTensorList* x,
                                  aclTensorList* weight,
                                  aclTensorList* scale,
                                  aclTensorList* perTokenScale,
                                  aclTensor* groupedList,
                                  aclTensorList* y,
                                  aclrtStream stream) {
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;

    auto ret = aclnnIndexGroupMatmulGetWorkspaceSize(x,
                                                     weight,
                                                     scale,
                                                     perTokenScale,
                                                     groupedList,
                                                     y,
                                                     &workspaceSize,
                                                     &executor);
    CHECK_RET(
        ret == ACL_SUCCESS,
        LOG_PRINT("aclnnGroupedMatmulGetWorkspaceSize failed. ERROR: %d\n",
                  ret);
        return ret);

    void* workspaceAddr = nullptr;
    if (workspaceSize > 0) {
      ret =
          aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
      CHECK_RET(ret == ACL_SUCCESS,
                LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret);
                return ret);
    }

    ret = aclnnIndexGroupMatmul(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclnnIndexGroupMatmul failed. ERROR: %d\n", ret);
              return ret);

    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
              return ret);

    if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
    }

    return ACL_SUCCESS;
  }
  aclTensor* groupedList = nullptr;
  aclTensorList* perTokenScale = nullptr;
  aclTensorList* x = nullptr;
  aclTensorList* weight = nullptr;
  aclTensorList* scale = nullptr;
  aclTensorList* y_index = nullptr;
  GroupGemmBase base;
};

class GroupGemmNative {
 public:
  GroupGemmNative(GroupGemmBase& base) : base(base) {}
  void process(aclrtStream stream) {
    this->createTensors();
    CHECK_ACL_SUCCESS(
        this->execute_group_gemm_native_operator(
            x, weight, scale, perTokenScale, groupedList, y_grouped, stream),
        "executeMatmulOperator failed");
  }
  void destroyTensors() {
    GroupGemmNative::clean_up(x, weight, scale, perTokenScale, y_grouped);
  }

 private:
  void createTensors() {
    CHECK_ACL_SUCCESS(utils::create_tensor_list_from_torch(
                          this->base.shapes.x_shape[0], this->base.input_a, &x),
                      "create x TensorList failed");
    CHECK_ACL_SUCCESS(
        utils::create_tensor_list_from_torch(
            this->base.shapes.scale_shape[0], this->base.scales, &scale),
        "create scale TensorList failed");
    CHECK_ACL_SUCCESS(utils::create_tensor_list_from_torch(
                          this->base.shapes.per_token_scale_shape[0],
                          this->base.per_token_scales,
                          &perTokenScale),
                      "create perTokenScale TensorList failed");
    CHECK_ACL_SUCCESS(
        utils::create_tensor_from_torch(this->base.shapes.group_list_shape[0],
                                        this->base.group_offset,
                                        &groupedList),
        "create groupedList Tensor failed");
    CHECK_ACL_SUCCESS(
        this->base.process_weight_tensor(
            this->base.shapes.weight_shape, this->base.input_b, &weight),
        "process weight tensor failed");
    CHECK_ACL_SUCCESS(
        utils::create_tensor_list_from_torch(this->base.shapes.y_shape[0],
                                             this->base.output_grouped,
                                             &y_grouped),
        "create y_grouped TensorList failed");
  }
  static void clean_up(aclTensorList* x,
                       aclTensorList* weight,
                       aclTensorList* scale,
                       aclTensorList* perTokenScale,
                       aclTensorList* y) {
    aclDestroyTensorList(x);
    aclDestroyTensorList(weight);
    aclDestroyTensorList(scale);
    aclDestroyTensorList(perTokenScale);
    aclDestroyTensorList(y);
  }
  int execute_group_gemm_native_operator(aclTensorList* x,
                                         aclTensorList* weight,
                                         aclTensorList* scale,
                                         aclTensorList* perTokenScale,
                                         aclTensor* groupedList,
                                         aclTensorList* y,
                                         aclrtStream stream) {
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    aclTensorList* bias = nullptr;
    aclTensorList* offset = nullptr;
    aclTensorList* antiquantScale = nullptr;
    aclTensorList* antiquantOffset = nullptr;
    aclTensorList* activationInput = nullptr;
    aclTensorList* activationQuantScale = nullptr;
    aclTensorList* activationQuantOffset = nullptr;
    aclTensorList* activationFeatureOut = nullptr;
    aclTensorList* dynQuantScaleOut = nullptr;
    int64_t splitItem = 2;
    int64_t groupType = 0;
    int64_t groupListType = 0;
    int64_t actType = 0;

    auto ret = aclnnGroupedMatmulV4GetWorkspaceSize(x,
                                                    weight,
                                                    bias,
                                                    scale,
                                                    offset,
                                                    antiquantScale,
                                                    antiquantOffset,
                                                    perTokenScale,
                                                    groupedList,
                                                    activationInput,
                                                    activationQuantScale,
                                                    activationQuantOffset,
                                                    splitItem,
                                                    groupType,
                                                    groupListType,
                                                    actType,
                                                    y,
                                                    activationFeatureOut,
                                                    dynQuantScaleOut,
                                                    &workspaceSize,
                                                    &executor);
    CHECK_RET(
        ret == ACL_SUCCESS,
        LOG_PRINT("aclnnGroupedMatmulGetWorkspaceSize failed. ERROR: %d\n",
                  ret);
        return ret);

    void* workspaceAddr = nullptr;
    if (workspaceSize > 0) {
      ret =
          aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
      CHECK_RET(ret == ACL_SUCCESS,
                LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret);
                return ret);
    }

    ret = aclnnGroupedMatmulV4(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclnnIndexGroupMatmul failed. ERROR: %d\n", ret);
              return ret);

    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
              return ret);

    if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
    }

    return ACL_SUCCESS;
  }
  aclTensor* groupedList = nullptr;
  aclTensorList* perTokenScale = nullptr;
  aclTensorList* x = nullptr;
  aclTensorList* weight = nullptr;
  aclTensorList* scale = nullptr;
  aclTensorList* y_grouped = nullptr;
  GroupGemmBase base;
};

}  // namespace group_gemm

#endif